import torch
import itertools

class Iter1:
    def fun(self):
        while True:
            yield from torch.arange(100).tolist()
    
    def fun2(self):
        g = torch.Generator()
        g.manual_seed(1236)
        while True:
            listss = torch.randperm(5, generator=g).tolist()
            outls = []
            for vvi in listss:
                outls.append(vvi*2)
                outls.append(vvi*2 + 1)
            print(outls)
            yield from outls
    def __iter__(self):
        # yield from itertools.islice(self.fun(), 0, None, 2)
        rank = 1
        batch_size = 6
        world_size = 4
        idx = 0
        for idvv in itertools.islice(self.fun2(), 0, None, 1):
            idx += 1
            yield idx,idvv      
                      
    def s__iter__(self):
        # yield from itertools.islice(self.fun(), 0, None, 2)
        rank = 1
        batch_size = 6
        world_size = 4
        for idx in itertools.islice(self.fun2(), 0, None, 1):
            v1 = idx // batch_size
            v2 = v1 % world_size
            if v2 == rank:
                yield idx


def ffun():
    while True:
        yield from torch.arange(10).tolist()


def ffiter():
    yield from itertools.islice(ffun(), 0, None, 2)
pp = ffiter

iterobj = Iter1()
def ssam():
    yield from iterobj

sam = ssam()
for s in sam:
    print(s)

p = Iter
p = list(res)
print(','.join(p))
res = 2